from collections import deque
from typing import Any, SupportsFloat
import gymnasium as gym
import numpy as np

class RecordSuccessRateWrapper(gym.Wrapper):
    def __init__(self, env: gym.Env):
        super().__init__(env)
        self.recorder = deque(maxlen=20)
        self.difficulty_level = 0 #由于浮点数的计算问题 -0.1+0.1!=0，我又不想改环境，在这里新加一个变量来记录难度

    def step(self, action: Any):
        observation, reward, terminated ,truncated ,info = self.env.step(action)
        if terminated:
            self.recorder.append(1)
            info = self.update_info(info)
        elif truncated:
            self.recorder.append(0)
            info = self.update_info(info)         
        return observation, reward, terminated, truncated, info
    
    def update_info(self, info):
        if len(self.recorder)==self.recorder.maxlen:
            info['success_rate'] = sum(self.recorder) / len(self.recorder)
            self.env.unwrapped.success_rate = info['success_rate']
            if self.env.unwrapped.curriculum :
                if self.difficulty_level < 10 and self.env.unwrapped.success_rate>=0.8:
                    self.difficulty_level += 1
                    self.env.unwrapped.curriculum_difficulty = self.difficulty_level/10
                    for _ in range(15):
                        self.recorder.popleft()

                elif self.difficulty_level > 0 and self.env.unwrapped.success_rate<=0.5:
                    self.difficulty_level -= 1
                    self.env.unwrapped.curriculum_difficulty = self.difficulty_level / 10
                    for _ in range(15):
                        self.recorder.popleft()

        return info

    